package com.chatplus.application.aiprocessor.channel.chat.chatgpt;

import cn.hutool.json.JSONObject;
import com.baomidou.mybatisplus.core.toolkit.Wrappers;
import com.chatplus.application.aiprocessor.listener.impl.ChatGPTFunctionEventSourceListener;
import com.chatplus.application.aiprocessor.provider.ChatAiProcessorServiceProvider;
import com.chatplus.application.aiprocessor.channel.chat.ChatAiProcessorService;
import com.chatplus.application.aiprocessor.enumeration.MessageTypeEnum;
import com.chatplus.application.aiprocessor.function.TriggerFunctionService;
import com.chatplus.application.aiprocessor.handler.dto.ChatRecordMessage;
import com.chatplus.application.aiprocessor.listener.impl.ChatGPTEventSourceListener;
import com.chatplus.application.aiprocessor.provider.TriggerFunctionServiceProvider;
import com.chatplus.application.aiprocessor.util.ChatAiUtilHandler;
import com.chatplus.application.common.logging.SouthernQuietLogger;
import com.chatplus.application.common.logging.SouthernQuietLoggerFactory;
import com.chatplus.application.domain.dto.FunctionParametersDto;
import com.chatplus.application.domain.entity.functions.FunctionEntity;
import com.chatplus.application.enumeration.AiPlatformEnum;
import com.chatplus.application.service.functions.FunctionService;
import com.google.common.collect.Lists;
import com.unfbx.chatgpt.OpenAiClient;
import com.unfbx.chatgpt.OpenAiStreamClient;
import com.unfbx.chatgpt.entity.chat.*;
import com.unfbx.chatgpt.function.KeyRandomStrategy;
import com.unfbx.chatgpt.interceptor.OpenAILogger;
import okhttp3.OkHttpClient;
import okhttp3.logging.HttpLoggingInterceptor;
import org.apache.commons.collections4.CollectionUtils;
import org.apache.commons.lang3.StringUtils;
import org.springframework.stereotype.Service;

import java.util.*;
import java.util.concurrent.CountDownLatch;
import java.util.concurrent.TimeUnit;
import java.util.regex.Matcher;
import java.util.regex.Pattern;
import java.util.stream.Collectors;

/**
 * ChatGPT 机器人处理器
 *
 * @author chj
 * @date 2023/12/27
 **/
@Service(value = ChatAiProcessorServiceProvider.SERVICE_NAME_PRE + "OPENAI")
public class ChatGPTProcessor extends ChatAiProcessorService {
    private static final SouthernQuietLogger LOGGER = SouthernQuietLoggerFactory.getLogger(ChatGPTProcessor.class);
    private final FunctionService functionService;
    private final TriggerFunctionServiceProvider triggerFunctionServiceProvider;
    protected ChatAiUtilHandler chatAiUtilHandler;

    public ChatGPTProcessor(FunctionService functionService, TriggerFunctionServiceProvider triggerFunctionServiceProvider) {
        this.functionService = functionService;
        this.triggerFunctionServiceProvider = triggerFunctionServiceProvider;

    }

    @Override
    public void processStream(List<ChatRecordMessage> messageContext) throws Exception {
        initConfig();
        List<Message> messages = handleMessage(messageContext);
        String modelName = getSession().getAttributes().get("modelName").toString();
        // 以下为发起请求
        String prompt = messageContext.stream()
                .filter(msg -> msg.getReply() == null)
                .map(ChatRecordMessage::getPrompt)
                .collect(Collectors.joining());
        chatAiUtilHandler = new ChatAiUtilHandler(getSession(), getChannel(), prompt);
        sendOpenAi(messages, modelName, prompt);
    }

    @Override
    public String processSync(List<ChatRecordMessage> messageContext) {
        initConfig();
        StringBuilder prompt = new StringBuilder();
        OpenAiClient openAiClient = OpenAiClient.builder()
                .apiHost(handleHost(getApiUrl()))
                .apiKey(getApiKeyList())
                .keyStrategy(new KeyRandomStrategy())
                .okHttpClient(getOkHttpClient())
                .build();
        List<Message> messages = handleMessage(messageContext);
        //聊天模型：gpt-3.5
        ChatCompletion chatCompletion = ChatCompletion.builder().messages(messages).build();
        ChatCompletionResponse chatCompletionResponse = openAiClient.chatCompletion(chatCompletion);
        chatCompletionResponse.getChoices().forEach(e -> prompt.append(e.getMessage().getContent()));
        return prompt.toString();
    }

    @Override
    public AiPlatformEnum getChannel() {
        return AiPlatformEnum.OPEN_AI;
    }

    @Override
    public void stopChat(String sessionId) {

    }

    private void sendOpenAi(List<Message> messages, String modelName, String prompt) throws Exception {
        OpenAiStreamClient openAiStreamClient = OpenAiStreamClient.builder()
                .apiHost(handleHost(getApiUrl()))
                .apiKey(getApiKeyList())
                //自定义key使用策略 默认随机策略
                .keyStrategy(new KeyRandomStrategy())
                .okHttpClient(getOkHttpClient())
                .build();
        List<Functions> functionsList = chatFunction();
        ChatCompletion chatCompletion = ChatCompletion.builder()
                .model(modelName)
                .temperature(getCommonSetting().getTemperature())
                .maxTokens(getCommonSetting().getMaxTokens())
                .messages(messages)
                .stream(true)
                .functions(functionsList)
                .functionCall("auto")
                .build();
        CountDownLatch countDownLatch = new CountDownLatch(1);
        //ChatGPTEventSourceListener eventSourceListener = new ChatGPTEventSourceListener(getSession(), getChannel(), prompt,countDownLatch);
        ChatGPTEventSourceListener eventSourceListener = new ChatGPTEventSourceListener(getSession(), getChannel(), prompt,countDownLatch);
        openAiStreamClient.streamChatCompletion(chatCompletion, eventSourceListener);
        countDownLatch.await();
        String args = eventSourceListener.getArguments();
        String name = eventSourceListener.getName();
        functionAction(name, args);
    }
    /**
     * 函数设置
     *
     * @return 返回函数列表
     */
    public List<Functions> chatFunction() {
        List<FunctionEntity> functionList = functionService.list(Wrappers.<FunctionEntity>lambdaQuery().eq(FunctionEntity::getEnabled, Boolean.TRUE));
        if (CollectionUtils.isEmpty(functionList)) {
            return Lists.newArrayList();
        }
        return functionList.stream().map(functionEntity -> {
            FunctionParametersDto functionParametersDto = functionEntity.getParameters();
            Map<String, FunctionParametersDto.FunctionFieldDetailBean> propertiesMap
                    = functionParametersDto.getProperties();
            JSONObject properties = new JSONObject();
            if (propertiesMap != null && !propertiesMap.isEmpty()) {
                propertiesMap.forEach((k, v) -> {
                    JSONObject propertiesField = new JSONObject();
                    propertiesField.putOpt("type", v.getType());
                    propertiesField.putOpt("description", v.getDescription());
                    properties.putOpt(k, propertiesField);
                });
            }
            Parameters parameters = Parameters.builder()
                    .type("object")
                    .properties(properties)
                    .required(functionParametersDto.getRequired()).build();
            return Functions.builder()
                    .name(functionEntity.getName())
                    .description(functionEntity.getDescription())
                    .parameters(parameters)
                    .build();
        }).toList();
    }

    public void functionAction(String name, String arguments) {
        if (StringUtils.isEmpty(name)) {
            return;
        }
        LOGGER.message("调用函数工具").context("name", name).context("arguments", arguments).info();
        FunctionEntity functionEntity = functionService.getOne(Wrappers.<FunctionEntity>lambdaQuery().eq(FunctionEntity::getName, name));
        if (Objects.isNull(functionEntity)) {
            return;
        }
        try {
            sendMessage(MessageTypeEnum.WS_START.getValue());
            // "正在调用工具 `%s` 作答 ...\n\n"
            String prompt = String.format("正在调用工具 `%s` 作答 ...%n%n", functionEntity.getLabel());
            sendMessage(prompt);
            TriggerFunctionService triggerFunctionService = triggerFunctionServiceProvider.getTriggerFunctionService(name);
            String body = triggerFunctionService.executeReturnMarkdown(functionEntity, arguments);
            if (StringUtils.isEmpty(body)) {
                sendMessage("调用函数工具出错：返回结果为空");
                return;
            }
            sendMessage(body);
            chatAiUtilHandler.saveChatHistory(body, 0L);
        } catch (Exception e) {
            sendMessage("调用函数工具出错，请等待管理员修复再重试");
            LOGGER.message("调用函数工具出错")
                    .context("name", name)
                    .context("arguments", arguments)
                    .exception(e)
                    .error();
        } finally {
            sendMessage(MessageTypeEnum.WS_END.getValue());
        }
    }

    private String handleHost(String url) {
        Pattern pattern = Pattern.compile("(https?://[^/]+/)");
        Matcher matcher = pattern.matcher(url);
        if (matcher.find()) {
            return matcher.group(1);
        }
        return null;
    }

    private OkHttpClient getOkHttpClient() {
        HttpLoggingInterceptor httpLoggingInterceptor = new HttpLoggingInterceptor(new OpenAILogger());
        //!!!!!!测试或者发布到服务器千万不要配置Level == BODY!!!!
        httpLoggingInterceptor.setLevel(HttpLoggingInterceptor.Level.HEADERS);
        return new OkHttpClient.Builder()
                .addInterceptor(httpLoggingInterceptor)
                .connectTimeout(30, TimeUnit.SECONDS)
                .writeTimeout(600, TimeUnit.SECONDS)
                .readTimeout(600, TimeUnit.SECONDS).build();
    }

    private List<Message> handleMessage(List<ChatRecordMessage> messageContext) {
        List<Message> messages = new ArrayList<>();
        if (CollectionUtils.isNotEmpty(messageContext)) {
            messages = messageContext.stream().map(msg -> Message.builder().content(msg.getPrompt()).role(BaseMessage.Role.USER).build()).collect(Collectors.toList());
        } else {
            Message currentMessage = Message.builder().content("你好").role(BaseMessage.Role.USER).build();
            messages.add(currentMessage);
        }
        return messages;
    }

}
